'''
Created on Feb 2, 2010

@author: oabalbin
'''

#import sys;
#import string;
#import StringIO;
import os;
import sys;
import glob;
import MySQLdb;
import numpy as np
from optparse import OptionParser


# My libraries
import signatures.parsers.read_microarray_data as pm
import signatures.parsers.read_gene_lists as gpm
import signatures.db.tables as dt
import signatures.db.query as dq
import signatures.common.classes as cm #record, arrayMatrix
import signatures.preprocess.tools as tl #filter_constant_genes, get_variable_genes, 
import signatures.seedgenes.seedlist as sd #get_seed_genes, write_varin_file
import signatures.bfrm.paramfile as bp
import signatures.bfrm.bfrm_run as br


class pipeline():
    
    def __init__(self, dbhost, dbuser, dbpasswd, database):
        
        self.mt = dt.tables(dbhost, dbuser, dbpasswd, database)
        self.mq = dq.query(dbhost, dbuser, dbpasswd, database)
        #table information
        #NSC,CONCUNIT,LCONC,PANEL,CELL,PANELNBR,CELLNBR,NLOGGI50,INDN,TOTN,STDDEV        
        #17 ,M ,-4.00,Non-Small Cell Lung ,NCI-H23 ,1 ,1 ,4.021,2 ,2 ,0.029
        
        self.tableDescription = '''NSC_ID INT, CONCUNIT VARCHAR(40), LCONC FLOAT(20,14), PANEL VARCHAR(60), CELL VARCHAR(40), 
                                 PANELNBR SMALLINT, CELLNBR SMALLINT, NLOGGI50 FLOAT (20,14), INDN SMALLINT,TOTN SMALLINT, STDDEV FLOAT(20,14), 
                                 INDEX NSC_ID (NSC_ID), INDEX CELL (CELL), INDEX PANEL (PANEL)'''
        self.tableName = 'Drugs_GI50'
        
        # Parameters for pre-processing data:
        # percentage of low variable gene to be filter out of the analysis
        self.nmasked = 40
        self.nmeas = 15
    
    def parse_input_annotfiles(self, annotInputFile):
    
        expr_file, seed_file, outdirs, db_file, samples, bfrm, typefile={},{},{},{},{},{},'nofile'
    
        for line in annotInputFile:
            fields = line.strip('\n')
                 
            if fields[0][0] == '':
                continue
            
            if fields[0][0] == '#':
                typefile = fields
                continue
            
            fields = fields.split('=')
            dict = {'#Array expression':expr_file, '#Seed genes':seed_file,'#Database file':db_file,'#Samples':samples,'#Output Directories':outdirs,\
                    '#BFRM':bfrm}[typefile]
            dict[fields[0]] = fields[1]
            
        return expr_file, seed_file, db_file, samples, bfrm, outdirs
    

    def read_files_folder(self,folderpath,ext):
        ''' '''
        # Read files in folder
        myfiles=[]
        for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
            myfiles.append(infile)
            #print "current file is: " + infile
        return myfiles
    
    
    def process_microarray_files(self,exp_files,tp):
        dbf=[]
        for f in exp_files:
            of=f+'_dbfile'
            tp.microarray_nci60_file(open(f),open(of,'w'))
            dbf.append(of)
        
        return dbf
    
    def create_database(self,dbfiles,doit=False):
        """
        It creates a databases if doit = True
        """
        if doit:
            print 'Making tables'            
            for dbf in dbfiles:     
                if not(self.mt.table_existence(self.tableName)):
                    self.mt.create_table(self.tableName,self.tableDescription)
                    
                self.mt.load_data_table_cvs(self.tableName,dbf)
    
    
    
    def write_compound_mat(self, expMat, vargenes, matfile, genefile):
        """
        write an compound dataset. 
        """
        
        genelist = np.array(expMat.geneSym)
        genelist = genelist[vargenes]
        
        vg_expMat = expMat.expVal[vargenes,:]
        
        samples = [ i[1] for i in expMat.sample]
        matfile.write(",".join(samples).replace(",", "\t")+'\n')
        for i, g in enumerate(genelist):
            el = map( str,list(vg_expMat[i,:]) ) 
            #genefile.write(g+'\n')
            matfile.write(g+"\t"+",".join(list( el )).replace(",", "\t") + '\n')
    
    def pickledump_compound_mat(self, expMat, varcpm, matfile, genefile):
        """
        Pickle the compund data set, so we can used for later analysis as 
        the correlation with factor activity.
        varcpm allows to subset the data
        """
        
        genelist = np.array(expMat.geneSym)
        genelist = genelist[varcpm]
        probeID = np.array(expMat.probeID)
        probeID = probeID[varcpm]
        vg_expMat = expMat.expVal[varcpm,:]
        matdim = vg_expMat.shape
        
        samples = [ i[1] for i in expMat.sample]
        genes = [g for g in genelist]
        genefile.write(",".join(samples).replace(",", "\t")+'\n')
        genefile.write(",".join(genes).replace(",", "\t")+'\n')
        
        #### Copying the array matrix
        odexpMat=cm.arrayMatrix(matdim[0], matdim[1])
        
        # Making a copy of the old expMat fields    
        odexpMat.probeID = list(probeID)
        odexpMat.geneSym = list(genelist)
        odexpMat.sample = expMat.sample
        odexpMat.geneSynd = set()
        odexpMat.genInd = {}
        # The new matrix
        odexpMat.expVal = vg_expMat
        
        print 'The shape of the matrix to dump is '
        print vg_expMat.shape
        
        np.ma.dump(odexpMat, matfile)
    
    
    def get_cmp_toprint(self, expMat, varcpm, nmeas, nmasked):
        """
        It gives the compounds present in more than n samples.
        """
        cmppr=[]
        vg_expMat = expMat.expVal[varcpm,:]
        
        for i in varcpm:
            # Compounds for which the value for all cell lines was 4.0. 
            # I think that is a error in measure. NOTE: Need to check with the Drug website
            unm = len(vg_expMat[i, vg_expMat[i,:] == 4.0])
            # Number of non-masked values:
            nmad = np.ma.count(vg_expMat[i,:])
            
            if (unm <= nmeas and nmad >= nmasked):
                cmppr.append(i)
            
        
        print  len(cmppr)
        
        return cmppr     
             
             
##########

if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--annotFile", dest="annotFile",
                            help="annotation file for all files to use")
    optionparser.add_option("-d", action="store_true", dest="dbtrue",
                            help="create a database. Bool")
    optionparser.add_option("-u", action="store_true", dest="matdump",
                            help="create a dumpfile. Bool")
    optionparser.add_option("-l", action="store_true", dest="matload",
                            help="load a expMat pickle file. Bool")


    
    (options, args) = optionparser.parse_args()
       
    # start the pipeline. 
    # open the database
    pipe = pipeline("localhost", "oabalbin", "oscar", "signatures")
    expr_folder, seed_file, db_file, samples, bfrm, outdirs = pipe.parse_input_annotfiles(open(options.annotFile))    
    #read folder for expression profiles
    dbfiles = pipe.read_files_folder(expr_folder['expdir'],'.txt')

    ####### output files
    
    outfile1 = outdirs['outdir']+'drugsensitivy_mat.txt'  # sensitivity values
    outfile2 = outdirs['outdir']+'compound_cmp.txt'
    outfile3 = outdirs['outdir']+'drugsenst_mat_dump.txt'
    ######### Get raw data  ############
    # Get the database files
    tp = pm.parser()
    #dbfiles = pipe.process_microarray_files(exp_files,tp)
    
    ######## Create database
    pipe.mt.create_connector()
    pipe.create_database(dbfiles,options.dbtrue)
    
    ########## Make queries ############
    
    pipe.mq.create_connector()
    # list of samples:
    samplelist = tp.list_of_names_in_line(open(samples['samples']))
    # It gets any list of  cpm, samples or any field in the tableName
    compoundlist = pipe.mq.get_compound_list(pipe.tableName,'NSC_ID')    
    print 'The number of compounds to analyze  is ' + str(len(compoundlist))
    
    # Do query to obtain the GI50 matrix related with the gene factors
    expMat = pipe.mq.get_exp_mat_comp(samplelist, compoundlist, pipe.tableName)
    
    # pickle the expMat file
    if options.matdump:
        expMat = pipe.mq.get_exp_mat_comp(samplelist, compoundlist, pipe.tableName)
        np.ma.dump(expMat, '/home/oabalbin/projects/networks/output/drugsensitivity/drugsenst_matorig_dump.txt')
    if options.matload:
        expMat = np.ma.load('/home/oabalbin/projects/networks/output/drugsensitivity/drugsenst_matorig_dump.txt')
        
    # I want all the compounds recorded in the database.
    cmp_toprint =  pipe.get_cmp_toprint(expMat, range(len(compoundlist)), pipe.nmeas, pipe.nmasked)
    
    print 'The filtered number of compounds to analyze  is ' + str(len(cmp_toprint))
    
    pipe.write_compound_mat(expMat, cmp_toprint, open(outfile1,'w'), open(outfile2,'w')) 
    pipe.pickledump_compound_mat(expMat, cmp_toprint, open(outfile3,'w'), open(outfile2,'w'))
    
    
    
    
    
    sys.exit(0)